package com.mayee.me.gateway.filter;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.ArrayUtils;
import org.reactivestreams.Publisher;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.ResponseCookie;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.util.MultiValueMap;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Objects;
import java.util.zip.GZIPInputStream;
import java.util.zip.GZIPOutputStream;

/**
 * @program: micro-service
 * @description:
 * @author: Bobby.Ma
 * @create: 2021-01-09 09:03
 **/
@Slf4j
public class ResponseDecorator extends ServerHttpResponseDecorator {

    private byte[] bytes;

    public ResponseDecorator(ServerHttpResponse delegate) {
        super(delegate);
    }

    public ResponseDecorator(ServerHttpResponse delegate,byte[] bytes) {
        super(delegate);
        this.bytes = bytes;
    }

    @Override
    public MultiValueMap<String, ResponseCookie> getCookies() {
        return super.getCookies();
    }

    @Override
    public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
        Flux<? extends DataBuffer> fluxBody = Flux.from(body);
        ServerHttpResponse response = getDelegate();
        DataBufferFactory bufferFactory = response.bufferFactory();
        if (Objects.nonNull(bytes)){
            return super.writeWith(Flux.just(bufferFactory.wrap(bytes)));
        }
        Flux<DataBuffer> flux = fluxBody.buffer().map(dataBuffers -> {
            DataBuffer join = bufferFactory.join(dataBuffers);
            byte[] content = new byte[join.readableByteCount()];
            join.read(content);
            // 释放掉内存
            DataBufferUtils.release(join);
            response.getHeaders().setContentLength(content.length);
            //判断响应是否有压缩
            boolean gzip = Objects.equals("gzip", response.getHeaders().getFirst("Content-Encoding"));
            String str = gzip ? new String(uncompresss(content), StandardCharsets.UTF_8) : new String(content);
            // todo 加密操作
            response.getHeaders().setContentLength(gzip ? compress(str.getBytes()).length : str.getBytes().length);
            return bufferFactory.wrap(str.getBytes(StandardCharsets.UTF_8));
        });
        return super.writeWith(flux);
    }

    /**
     * @param bytes
     * @Description: 解压字节
     * @return: byte[]
     * @Author: Bobby.Ma
     * @Date: 2020/12/28 16:33
     */
    private byte[] uncompresss(byte[] bytes) {
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        ByteArrayInputStream in = new ByteArrayInputStream(bytes);
        try (GZIPInputStream gzip = new GZIPInputStream(in)) {
            byte[] buffer = new byte[1024];
            int n;
            while ((n = gzip.read(buffer)) >= 0) {
                out.write(buffer, 0, n);
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
        return out.toByteArray();
    }

    /**
     * @param bytes
     * @Description: 压缩字节
     * @return: byte[]
     * @Author: Bobby.Ma
     * @Date: 2020/12/28 16:37
     */
    private byte[] compress(byte[] bytes) {
        if (ArrayUtils.isEmpty(bytes)) {
            return null;
        }
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        try (GZIPOutputStream gzip = new GZIPOutputStream(out)) {
            gzip.write(bytes);
        } catch (IOException e) {
            e.printStackTrace();
        }
        return out.toByteArray();
    }
}
